#!/bin/bash

# Copyright 2021  Mobvoi Inc(Author: Di Wu, Binbin Zhang)
#                 NPU, ASLP Group (Author: Qijie Shao)

. ./path.sh || exit 1;

gpu_list='0,1,2,3,4,5,6,7' # 推理时可见的gpu
# shellcheck disable=SC2154
export CUDA_VISIBLE_DEVICES="${gpu_list}"
gpu_id=

test_data_dir=/home/work_nfs8/xlgeng/data/scp_test
test_sets="aishell1 aishell2 test_net test_meeting speechio_0 speechio_1 speechio_2 speechio_3 speechio_4"

decode_checkpoint=/home/work_nfs8/xlgeng/new_workspace/checkpoint/1_comformer_task/20220506_u2pp_conformer_exp_wenetspeech/final.pt
train_yaml=/home/work_nfs8/xlgeng/new_workspace/checkpoint/1_comformer_task/20220506_u2pp_conformer_exp_wenetspeech/train.yaml
dir=exp/u2pp_conformer
decode_modes="ctc_greedy_search ctc_prefix_beam_search attention attention_rescoring"


decoding_chunk_size=-1
ctc_weight=0.5
reverse_weight=0.0
blank_penalty=0.0
length_penalty=0.0
decode_batch=16

. tools/parse_options.sh || exit 0;
test_sets=$(echo "$test_sets" | sed 's/---/ /g')
echo "test_data_dir is ${test_data_dir}"
echo "test_sets is ${test_sets}"
echo "dir is ${dir}"
echo "decode_checkpoint is ${decode_checkpoint}"
echo "decode_modes is ${decode_modes}"
echo "decoding_chunk_size is ${decoding_chunk_size}"
echo "ctc_weight is ${ctc_weight}"
echo "reverse_weight is ${reverse_weight}"
echo "blank_penalty is ${blank_penalty}"
echo "length_penalty is ${length_penalty}"
echo "decode_batch is ${decode_batch}"
echo "gpu_id is ${gpu_id}"
echo "test_sets is ${test_sets}"

set -u
set -o pipefail


echo "开始 Test model"
# 如果要平均模型，在加入代码，此处就先不加了
base=$(basename $decode_checkpoint)
result_dir_root=$dir/${base}_chunk${decoding_chunk_size}_ctc${ctc_weight}_reverse${reverse_weight}_blankpenalty${blank_penalty}_lengthpenalty${length_penalty}
echo "result_dir_root is ${result_dir_root}"
for testset in ${test_sets}; do
{
  echo "耿雪龙：开始 Testing ${testset} on GPU ${gpu_id}"
  result_dir=${result_dir_root}/${testset}
  mkdir -p "${result_dir}"
  echo "result_dir is ${result_dir} for ${testset}"
  # #decode_modes如果不加双引号就是多个值，也是这里需要的，如果有双引号就是一个值
  python wenet/bin/recognize.py --gpu ${gpu_id} \
    --modes $decode_modes \
    --config $train_yaml \
    --data_type "raw" \
    --test_data $test_data_dir/"$testset"/data.list \
    --checkpoint "$decode_checkpoint" \
    --beam_size 10 \
    --batch_size $decode_batch \
    --blank_penalty $blank_penalty \
    --length_penalty $length_penalty \
    --ctc_weight $ctc_weight \
    --reverse_weight $reverse_weight \
    --result_dir "$result_dir" \
    ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}
  echo '推理完毕，开始计算wer'
  for mode in ${decode_modes}; do
    echo "计算wer for ${mode} for ${testset}"
    python tools/compute-wer.py --char=1 --v=1 \
      $test_data_dir/"$testset"/text "$result_dir"/"$mode"/text > "$result_dir"/"$mode"/wer
    echo "计算wer完毕 for ${mode} for ${testset}"
    tail -n 5 "$result_dir"/"$mode"/wer
  done
}
done
wait


# 导出模型
#if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
#  echo "Export the best model you want"
#  python wenet/bin/export_jit.py \
#    --config $dir/train.yaml \
#    --checkpoint $dir/avg_${average_num}.pt \
#    --output_file $dir/final.zip
#fi

#  平均模型
#  if [ ${average_checkpoint} == true ]; then
#    decode_checkpoint=$dir/avg_${average_num}.pt
#    echo "do model average and final checkpoint is $decode_checkpoint"
#    python wenet/bin/average_model.py \
#        --dst_model $decode_checkpoint \
#        --src_path $dir  \
#        --num ${average_num} \
#        --mode ${average_mode} \
#        --max_step ${max_step} \
#        --val_best
#  fi